#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""Cog Model"""
__author__ = "Lei Chen <hide1713@gmail.com>"
import logging
import copy
from genome_exception import GCogNotFindException
from cog import Cog
from function_category import FunctionCategory

class CogModel(object):
    """Cog model store the cog tree and hit information.
    We will use the data in CogModel to draw things.
    Each CogModel represent a data file.
    """
    
    def __init__(self, file_base_name, cog_map, owner=None):
        """One data file should have one CogModel 
        - `file_base_name`: is the file name with out path
        - `cog_map`: is the cog family map
        """
        self.log = logging.getLogger(self.__class__.__name__)
        self._cog_map = cog_map #Reference of cog families. Do not change it
        self._genome_data_map = {} #Restore the genome records
        self._base_name = file_base_name
        self.owner = owner #User id, the owner of this file

    def _get_base_name(self):
        return self._base_name
    
    base_name = property(_get_base_name)
    
    def add_genome_records(self, records):
        """Add genome records to _genome_data_map
        according to its family name
        
        Arguments:
        - `records`: a list of CogRecords
        """
        self.log.info("Adding genome records to data model")
        for r in records:
            if self._genome_data_map.has_key(r.cog_name):
                self._genome_data_map[r.cog_name].add_hits(r)
            elif self._cog_map.has_key(r.cog_name):
                old_cog = self._cog_map[r.cog_name]
                new_cog = copy.deepcopy(old_cog)
                new_cog.file_name = self._base_name
                self._genome_data_map[r.cog_name] = new_cog
                
                self._genome_data_map[r.cog_name].add_hits(r)
            else:
                self.log.error(r.cog_name + " is not in COG")
                raise GCogNotFindException(r.cog_name)

    def clean_data(self, ):
        """Remove data from cog map
        """
        self._genome_data_map = {}

    def get_single_cog(self, cog_name):
        """Return a cog object
        
        Arguments:
        - `cog_name`:
        """
        return self._genome_data_map[cog_name]

    def get_non_empty_cog_names(self, ):
        """Return a list of (cog_name, file_name, hit_count) 
        """
        name_list = [(i.name, self._base_name, i.number_of_hits)\
                     for i in self._genome_data_map.values()]
        return name_list
    
    def get_total_cog_names_and_hits(self, ):
        """Return a list of [hit_count]. The size of list is equal to
        the number of COG
        Note: a record will return even there's no hits on the cog.
        In this case,  hit_count would be zero
        """
        name_map = {}
        #Set hits of every cog to zero
        for cog in self._cog_map:
            name_map[cog] = 0
        #Set the hits value of current data set
        for i in self._genome_data_map.values():
            name_map[i.name] = i.number_of_hits
        #Sort the hits by cog name
        keys = name_map.keys()
        keys.sort()
        return [name_map[k] for k in keys]

    
    def get_cog_range(self, cog_name):
        """This is the interface to draw cog range graphic
        Return:
        the length of this cog and
        the a list of (id, start, end) positions of every genome record
        Arguments:
        - `cog_name`: the cog name string 
        """
        try:
            cog = self._genome_data_map[cog_name]
            return cog.length, cog.get_hits_range()
        except KeyError:
            raise GCogNotFindException(cog_name)

    def get_function_category(self, ):
        """Return a FunctionCategory object"""
        fc = FunctionCategory(self.base_name)
        for cog in self._genome_data_map.values():
            fc.add_cog(cog)
        return fc
